{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "sys.path.append(os.path.abspath('../'))\n",
    "from common import *\n",
    "import configs\n",
    "import matplotlib.gridspec as gridspec\n",
    "\n",
    "from constants import OUTPUT_DIR\n",
    "\n",
    "fig_name = 'shading_gradient'\n",
    "fig_dir = join(FIGURE_DIR, fig_name)\n",
    "\n",
    "scene_name = 'bunny'\n",
    "used_configs = [configs.OnlyShadingGrad(), configs.Warp(), configs.FiniteDifferences()]\n",
    "output_dir = join(OUTPUT_DIR, scene_name, 'no-tex-12-hq', used_configs[0].name)\n",
    "\n",
    "y_offset = -0.1\n",
    "fontsize = 12\n",
    "n_rows = 2\n",
    "n_cols = 3\n",
    "total_width = COLUMN_WIDTH\n",
    "aspect = n_rows / n_cols\n",
    "fig = plt.figure(1, figsize=(total_width, aspect * total_width), constrained_layout=False)\n",
    "\n",
    "gs = fig.add_gridspec(n_rows, n_cols, wspace=0.025, hspace=0.025)\n",
    "r = 15\n",
    "row = 0\n",
    "\n",
    "names = ['Ignoring discontinuities', 'Ours', 'Reference (FD)']\n",
    "\n",
    "# Show gradient images\n",
    "for col, (config, config_name_pretty) in enumerate(zip(used_configs, names)):\n",
    "    ax = fig.add_subplot(gs[row, col])\n",
    "    fn = f'{scene_name}_{config.name}_x.exr'\n",
    "    img = read_img(join(fig_dir, scene_name, fn), tonemap=False)\n",
    "    img = img[50:-50, 50:-50, :]\n",
    "    aspect = img.shape[0] / img.shape[1]\n",
    "    img = np.repeat(np.repeat(img, 4, axis=0), 4, axis=1)\n",
    "    ax.imshow(np.mean(img, axis=-1), cmap='coolwarm_r', vmin=-r, vmax=r, interpolation='none')\n",
    "    disable_ticks(ax)\n",
    "    if col == 0:\n",
    "        ax.set_ylabel(f'Gradient', labelpad=5, fontsize=fontsize)\n",
    "    txt = ax.set_title(config_name_pretty, fontsize=fontsize, va='top', y=1.07)\n",
    "\n",
    "# Display optimization results\n",
    "row += 1\n",
    "used_configs = [configs.OnlyShadingGrad(), configs.Warp()]\n",
    "for col, (config, config_name_pretty) in enumerate(zip(used_configs, names)):\n",
    "    ax = fig.add_subplot(gs[row, col])\n",
    "    fn = f'{config.name}_final.exr'\n",
    "    img = read_img(join(fig_dir, scene_name, fn), tonemap=True)\n",
    "    aspect = img.shape[0] / img.shape[1]\n",
    "    img = np.repeat(np.repeat(img, 4, axis=0), 4, axis=1)\n",
    "    ax.imshow(img, interpolation='none')\n",
    "    disable_ticks(ax)\n",
    "    if col == 0:\n",
    "        ax.set_ylabel(f'Optimization', labelpad=5, fontsize=fontsize)\n",
    "    txt = ax.set_title(config_name_pretty, fontsize=fontsize, va='top', y=-0.1)\n",
    "\n",
    "    ax = fig.add_subplot(gs[row, -1])\n",
    "    if row == n_rows - 1:\n",
    "        txt = ax.set_title('Reference images \\n (12 views in total)', fontsize=fontsize, y=y_offset, va='top')\n",
    "    ax.set_axis_off()\n",
    "\n",
    "    gs_insets = gridspec.GridSpecFromSubplotSpec(2, 2, subplot_spec=gs[row, -1], wspace=0.00, hspace=0.05)\n",
    "    ref_views = [0, 3, 7, 10]\n",
    "    for idx, ref_view in enumerate(ref_views):\n",
    "        r = idx // 2\n",
    "        c = idx % 2\n",
    "        ax = fig.add_subplot(gs_insets[r, c])\n",
    "        img = read_img(join(output_dir, f'ref-{ref_view:02d}.exr'))\n",
    "        ax.imshow(img, interpolation='none')\n",
    "        disable_ticks(ax)\n",
    "\n",
    "plt.margins(0, 0)\n",
    "# save_fig(fig_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('mi')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c8cc192903caa17681aa39d71202092ac11a526d37a1c4ad2948f13605924304"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
